import functools
import time

from loguru import logger


def timeit(func):
    """
    用于计算函数执行时间的装饰器
    :param func:
    :return:
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        run_time = end_time - start_time
        logger.debug(f"函数 {func.__name__} 运行时间为：{run_time:.6f}秒")
        return result

    return wrapper


def tensor_to_2dim(func):
    """
    将一维的tensor转成二维的tensor
    :param func:
    :return:
    """

    @functools.wraps(func)
    def wrapper(*args):
        dec = args[1]
        dec = dec.clone().unsqueeze(dim=0) if dec.dim() == 1 else dec
        return func(*(args[0], dec))

    return wrapper


def tensor_to_1dim(func):
    """
    将二维的tensor转成一维的tensor
    :param func:
    :return:
    """

    @functools.wraps(func)
    def wrapper(*args):
        dec = args[1]
        dec = dec.squeeze(dim=0) if dec.dim() == 2 else dec
        return func(*(args[0], dec))

    return wrapper
